#include "thread_pool.hpp"
#include <atomic>
#include <cassert>
#include <iostream>
#ifdef _WIN32
    #include <MMSystem.h>
    #pragma comment(lib, "winmm.lib")
#endif // _WIN32

ThreadPool thread_pool(128, 16);
std::atomic<int> sum {0};

void print()
{
    ThreadFunc::SleepFor(100);
    std::cout << "GetThreadId: " << ThreadFunc::GetThreadId() << std::endl;
    ++sum;
}

class TestA
{
public:
    TestA()
    {
        start = std::chrono::steady_clock::now();
    }

    ~TestA() {}

    void print()
    {
        ++sum;
        end = std::chrono::steady_clock::now();
        assert(std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() < 1002);
    }

    std::chrono::steady_clock::time_point start;
    std::chrono::steady_clock::time_point end;
};

void Test1()
{
    std::shared_ptr<ThreadPool::TaskToken> task_token = std::make_shared<ThreadPool::TaskToken>();
    std::vector<std::thread> threads;
    for (int i = 0; i < 8; ++i)
    {
        threads.emplace_back([&task_token]() {
            for (int i = 0; i < 10; ++i)
            {
                thread_pool.AddTask(task_token, []() {
                    print();
                });
            }
        });
    }

    ThreadFunc::SleepFor(550);
    for (uint32_t i = 0; i < 8; ++i)
    {
        if (threads[i].joinable())
        {
            threads[i].join();
        }
    }
    assert(sum == 80);
    assert(task_token.use_count() == 1);
}

void Test2()
{
    sum = 0;
    std::shared_ptr<ThreadPool::TaskToken> task_token = std::make_shared<ThreadPool::TaskToken>();
    std::vector<std::thread> threads;
    for (int i = 0; i < 8; ++i)
    {
        threads.emplace_back([&task_token]() {
            for (int i = 0; i < 10; ++i)
            {
                auto test_ptr = std::make_shared<TestA>();
                thread_pool.AddTask(
                    task_token,
                    [test_ptr]() {
                        test_ptr->print();
                    },
                    1000);
            }
        });
    }

    ThreadFunc::SleepFor(5500);
    for (uint32_t i = 0; i < 8; ++i)
    {
        if (threads[i].joinable())
        {
            threads[i].join();
        }
    }
    assert(sum == 80);
    assert(task_token.use_count() == 1);
}

int main()
{
#ifdef _WIN32
    SetConsoleCP(65001);
    SetConsoleOutputCP(65001);
    timeBeginPeriod(1); /**< 开启精度设置 */
#endif
    Test1();
    Test2();
#ifdef _WIN32
    timeEndPeriod(1); /**< 结束精度设置 */
#endif
}